﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

using System.Diagnostics.Contracts;

using SimpleIR;
using PInfoModel;

namespace TestGen
{
	public class HeapGenerate
	{
		private readonly IRBodyDef m_bdf;
		private readonly PIModel m_pim;
		private readonly List<TGenConstraint> m_constraints;
		private readonly List<Tuple<string, string>> m_fixedvars;
		private readonly List<string> m_paramargs;

		public HeapGenerate(IRBodyDef bdf, PIModel pim, List<TGenConstraint> constraints, List<Tuple<string, string>> fixedvars, List<string> paramargs)
		{
			this.m_bdf = bdf;
			this.m_pim = pim;
			this.m_constraints = constraints;
			this.m_fixedvars = fixedvars;
			this.m_paramargs = paramargs;
		}

		public String BuildConstructor(string oparams, string callstr, int ccid)
		{
			StringBuilder csb = new StringBuilder(s_indent + "//" + callstr + "\n");
			String hsb = this.BuildConstructorArgList(oparams, this.m_paramargs);

			csb.Append(s_consheader + ccid);
			csb.Append(hsb);
			csb.Append("\t\t{");

			//add additional fixed vars
			csb.Append("\n" + s_indent + "//Additional fixed assignments.\n");
			csb.Append(s_indent + "#region fixed assignments\n");
			foreach(Tuple<string, string> fxv in this.m_fixedvars)
				csb.Append(s_indent + "int " + fxv.Item1 + " = " + fxv.Item2 + ";\n");
			csb.Append(s_indent + "#endregion\n\n");

			//add additional asserts constraints
			csb.Append(s_indent + "//Additional inferred constraints.\n");
			csb.Append(s_indent + "#region parameter constraints\n");
			foreach(TGenConstraint cst in this.m_constraints)
				csb.Append(s_indent + s_assume + "(" + cst.GenerateConstraintOutputForCSharp() + ");\n");
			csb.Append(s_indent + "#endregion\n\n");

			//allocate regions
			csb.Append(s_indent + "//Alloc Regions\n");
			csb.Append(s_indent + "#region allocate regions\n");
			String rgns = this.BuildRegionAllocs();
			csb.Append(rgns);
			csb.Append(s_indent + "#endregion\n\n");

			//add ptrs
			csb.Append(s_indent + "//Fill Edges\n");
			csb.Append(s_indent + "#region assign pointers\n");
			String edges = this.BuildEdgeAllocs();
			csb.Append(edges);
			csb.Append(s_indent + "#endregion\n\n");

			//check my other asserts
			csb.Append(s_indent + "//Test Other Properties as Needed\n");
			csb.Append(s_indent + "#region additional assertions\n");
			//should all be ok by construction at this point in time
			csb.Append(s_indent + "#endregion\n\n");

			//fill in scalars
			csb.Append(s_indent + "//Fill Scalars\n");
			csb.Append(s_indent + "#region assign scalars\n");
			String scalarinit = this.GenInitScalars();
			csb.Append(scalarinit);
			csb.Append(s_indent + "#endregion\n\n");

			//check all invariant constracts
			csb.Append(s_indent + "//INVARIANT CHECKS\n");
			csb.Append(s_indent + "#region invariant checks\n");

			csb.Append("\t//SKIPPING INV CHECK RIGHT NOW!!!\n");

			/*
			foreach(HeapRegion hr in this.m_nodes)
			{
				foreach(TypeIdentifier ttid in hr.NodeTypes)
				{
					if(ttid.SpecialHandleTypeTag == TypeIdentifier.SpecialTypeCategory.Normal)
					{
						String frgnct = GenerateNameForRegionAndTypeCount(hr, ttid);
						String frgnname = GenerateNameForRegionAndTypeName(hr, ttid);
						csb.Append(s_indent + "for(int jj = 0; jj < " + frgnct + "; ++jj) " + "PexInvariant.CheckInvariant(" + frgnname + "[jj]);\n");
					}
				}
			}
			*/
 
			csb.Append(s_indent + "#endregion\n\n");

			//invoke test method
			csb.Append(s_indent + "//Test\n");
			csb.Append(s_indent + callstr + ";\n");

			csb.Append("\t\t}\n");
			return csb.ToString();
		}

		/////////////////////////////////////////////////////////
		//String constants we use to build constructor

		private const String s_consheader = "\t\t[PexMethod]\n\t\tpublic void HeapAutoGen";
		private const String s_indent = "\t\t\t";
		private const String s_indent2 = s_indent + "\t";
		private const String s_indent3 = s_indent2 + "\t";
		private const String s_indent4 = s_indent3 + "\t";
		private const String s_assert = "PexAssert.IsTrue";
		private const String s_assume = "PexAssume.IsTrue";

		/// <summary>
		/// Emit code to declare all of the regions that we will need.
		/// </summary>
		private string BuildRegionAllocs()
		{
			StringBuilder csb = new StringBuilder();

			foreach(PIMNode hr in this.m_pim.Nodes)
			{
				Contract.Assert(!(hr.LocationType is IRStructType), "Not handled yet.");

				if(hr.HomeLocFor != null)
					continue;

				string rgnsize = NameGen.GenerateNameForRegionCount(hr);
				string rgnname = NameGen.GenerateNameForRegion(hr);

				if(hr.LocationType is IRContainerType)
				{
					IRContainerType ctype = (IRContainerType)hr.LocationType;
					IRType contentstype = (ctype.ContainerContentsType is IRPtrType) ? ((IRPtrType)ctype.ContainerContentsType).PtrTargetType : ctype.ContainerContentsType;
					if(ctype.IsArrayContainer)
						csb.Append(s_indent + contentstype.TypeName + "[][] " + rgnname + " = new " + contentstype + "[" + rgnsize + "][];\n");
					else if(ctype.IsListContainer)
						csb.Append(s_indent + "List<" + contentstype.TypeName + ">[] " + rgnname + " = new List<" + contentstype + ">[" + rgnsize + "];\n");
					else
						throw new NotImplementedException("Region of Containers " + ctype.ToString());
				}
				else
				{
					IRType rgntype = (hr.LocationType is IRPtrType) ? ((IRPtrType)hr.LocationType).PtrTargetType : hr.LocationType;
					csb.Append(s_indent + rgntype.TypeName + "[] " + rgnname + " = new " + rgntype.TypeName + "[" + rgnsize + "];\n");
				}

				csb.Append(s_indent + "for(int i = 0; i < " + rgnsize + "; ++i)\n");
				GenConsStringFor(csb, hr, rgnname);
				csb.Append("\n");
			}

			return csb.ToString();
		}

		/// <summary>
		/// Generate code that will allocate objects for a region array.
		/// </summary>
		private void GenConsStringFor(StringBuilder sb, PIMNode hr, string subrgnname)
		{
			sb.Append(s_indent + "{");

			if(hr.LocationType is IRContainerType)
			{
				IRContainerType ctype = (IRContainerType)hr.LocationType;
				IRType contentstype = (ctype.ContainerContentsType is IRPtrType) ? ((IRPtrType)ctype.ContainerContentsType).PtrTargetType : ctype.ContainerContentsType;
				string ccard = NameGen.GenerateNameForRegionCardinality(hr);

				if(ctype.IsArrayContainer)
					sb.Append(" " + subrgnname + "[i] = new " + contentstype.TypeName + "[" + ccard + "]" + "; ");
				else if(ctype.IsListContainer)
				{
					sb.Append("\n" + s_indent2 + subrgnname + "[i] = new List<" + contentstype.TypeName + ">();\n");
					sb.Append(s_indent2 + "for(int j = 0; j < " + ccard + "; ++j) { " + subrgnname + "[i].Add(default(" + contentstype + ")); } ");
					sb.Append("\n" + s_indent);
				}
				else
					throw new NotImplementedException("Region of Containers " + ctype.ToString());
			}
			else
			{
				IRType rgntype = (hr.LocationType is IRPtrType) ? ((IRPtrType)hr.LocationType).PtrTargetType : hr.LocationType;
				sb.Append(" " + subrgnname + "[i] = PexInvariant.CreateInstance<" + rgntype + ">(); ");
			}

			sb.Append("}\n");
		}

		private static int s_mselectctr = 0;

		/// <summary>
		/// Build the code to assign pointers between fields in the various regions
		/// </summary>
		private string BuildEdgeAllocs()
		{
			StringBuilder esb = new StringBuilder();

			foreach(IRRootIdentifier rr in this.m_pim.NameToNodeMap.Keys)
			{
				PIMNode homen = this.m_pim.ResolveHomeForRoot(rr);
				IRType hometype = (homen.LocationType is IRPtrType) ? ((IRPtrType)homen.LocationType).PtrTargetType : homen.LocationType;

				if(!(homen.LocationType is IRPtrType))
					continue;

				string vnv = (rr.RootNameString().Equals("this")) ? ("t" + rr.RootNameString()) : rr.RootNameString();
				esb.Append(s_indent + hometype.TypeName + " " + vnv + " = null;\n");

				if(homen.OutEdges.Count != 0)
				{
					string vnc = "pc_" + vnv;
					string vncsub = "vnc_" + vnv;

					if(homen.OutEdges.Count == 1 && homen.OutEdges[0].IsNonNull)
					{
						PIMNode tn = this.m_pim.Nodes[homen.OutEdges[0].TargetAddress];
						string tnname = NameGen.GenerateNameForRegion(tn);
						if(tn.InEdges.Count == 1)
							esb.Append(s_indent + vnv + " = " + tnname + "[0];\n");
						else
						{
							esb.Append(s_indent3 + s_assume + "(" + tnname + ".Length != 0);\n");
							esb.Append(s_indent + "int " + vncsub + " = PexChoose.ValueFromRange(\"trgtobj\", 0, " + tnname + ".Length);\n");
							esb.Append(s_indent + vnv + " = " + tnname + "[" + vncsub + "];\n");
						}
					}
					else
					{
						bool maynull = homen.OutEdges.Any((PIMEdge ee) => !ee.IsNonNull);
						esb.Append(s_indent + "int " + vnc + " = PexChoose.ValueFromRange(\"trgtrgn\", 0, " + (homen.OutEdges.Count + (maynull ? 1 : 0)) + ");\n");

						esb.Append(s_indent + "switch(" + vnc + ")\n");
						esb.Append(s_indent + "{\n");
						for(int j = 0; j < homen.OutEdges.Count; ++j)
						{
							string tname = NameGen.GenerateNameForRegion(this.m_pim.Nodes[homen.OutEdges[j].TargetAddress]);

							esb.Append(s_indent2 + "case " + j + ":\n");
							esb.Append(s_indent3 + s_assume + "(" + tname + ".Length != 0);\n");
							esb.Append(s_indent3 + "int " + vncsub + j + " = PexChoose.ValueFromRange(\"trgtobj\", 0, " + tname + ".Length);\n");
							esb.Append(s_indent3 + vnv + " = " + tname + "[" + vncsub + j + "];\n");

							esb.Append(s_indent3 + "break;\n");
						}


						esb.Append(s_indent2 + "default:\n");
						if(maynull)
							esb.Append(s_indent3 + vnv + " = null;\n");
						esb.Append(s_indent3 + "break;\n");

						esb.Append(s_indent + "}\n");
					}
				}

				esb.Append("\n");
			}

			foreach(PIMNode hr in this.m_pim.Nodes)
			{
				if(hr.HomeLocFor != null)
					continue;

				foreach(IRBasicAccessPath bap in hr.LocationType.TypeDescription.AllTerminalAccessPaths)
				{
					Contract.Assert(bap.Path.Length == 1, "We aren't that cool yet.");

					IRMemberField mfid = bap.PathFirst;
					if(!(mfid.FieldType is IRPtrType))
						continue;

					List<PIMEdge> possibleedges = hr.OutEdges.Where((PIMEdge ee) => ee.OffsetInSource == bap).ToList();
					if(possibleedges.Count == 0)
						continue;

					bool nonnull = possibleedges.All((PIMEdge ee) => ee.IsNonNull);
					bool nonintf = possibleedges.All((PIMEdge ee) => ee.IsInjective);

					if(possibleedges.Count == 1 && this.m_pim.Nodes[possibleedges[0].TargetAddress].InEdges.Count == 1 && nonintf && nonnull)
					{
						PIMNode tn = this.m_pim.Nodes[possibleedges[0].TargetAddress];
						string sel = this.BuildRegionObjSuperSimpleSelection(NameGen.GenerateNameForRegion(hr),  NameGen.GenerateNameForRegion(tn), mfid);
						esb.Append(sel);
					}
					else
					{
						esb.Append("\n");

						List<string> possibletargetsets = new List<string>();
						foreach(PIMEdge ee in possibleedges)
							possibletargetsets.Add(NameGen.GenerateNameForRegion(this.m_pim.Nodes[ee.TargetAddress]));

						if(nonintf)
						{
							for(int i = 0; i < possibletargetsets.Count; ++i)
							{
								String nrgn = possibletargetsets[i] + s_mselectctr++;
								esb.Append(s_indent + "var " + nrgn + " = " + possibletargetsets[i] + ".ToArray();\n");
								possibletargetsets[i] = nrgn;
							}
						}

						String sel = this.BuildRegionObjGeneralSelection(NameGen.GenerateNameForRegion(hr), possibletargetsets, mfid, (!nonnull), (!nonintf));
						esb.Append(sel);
					}
				}
			}

			return esb.ToString();
		}

		private static int rgnidexctr = 0;

		/// <summary>
		/// Build code to do cross region pointer assignement in the simple case (one choice, nonnull, and injective).
		/// </summary>
		private string BuildRegionObjSuperSimpleSelection(string sourcergn, string trgtrgn, IRMemberField mfid)
		{
			StringBuilder sb = new StringBuilder();

			if(mfid.IsSummaryField)
			{
				string nrgctr = "i_" + rgnidexctr++;
				sb.Append(s_indent + "int " + nrgctr + " = 0;\n");
				sb.Append(s_indent + "for(int sio = 0; sio < " + sourcergn + ".Length; ++sio)\n");
				sb.Append(s_indent + "{\n");
				sb.Append(s_indent2 + s_assume + "(" + nrgctr + " < " + trgtrgn + ".Length);\n");
				sb.Append(s_indent2 + "for(int sim = 0; sim < " + sourcergn + "[sio]." + (((IRContainerType)mfid.DefinedInClass).IsArrayContainer ? "Length" : "Count") + "; ++sim)");
				sb.Append("{ " + sourcergn + "[sio][sim] = " + trgtrgn + "[" + nrgctr + "++]; }\n");
				sb.Append(s_indent + "}\n");
			}
			else
			{
				string mfname = mfid.FieldName;
				IRType ftype = (mfid.FieldType is IRPtrType ? ((IRPtrType)mfid.FieldType).PtrTargetType : mfid.FieldType);

				sb.Append(s_indent + "for(int k = 0; k < " + sourcergn + ".Length; ++k) { ");
				sb.Append("PexInvariant.SetField<" + ftype.ToString() + ">(");
				sb.Append(sourcergn + "[k]," + " \"" + mfname + "\", " + trgtrgn + "[k]); }\n");
			}

			return sb.ToString();
		}

		/// <summary>
		/// Build code to do cross region pointer assignement in the complex case (multi-choice, may-null, or may-alias).
		/// </summary>
		private String BuildRegionObjGeneralSelection(string sourcergn, List<string> trgtrgns, IRMemberField mfid, bool maynull, bool mayintf)
		{
			trgtrgns.Sort();
			StringBuilder sb = new StringBuilder();

			if(mfid.IsSummaryField)
			{
				string cards = ((IRContainerType)mfid.DefinedInClass).IsArrayContainer ? "Length" : "Count";

				sb.Append(s_indent + "for(int sio = 0; sio < " + sourcergn + ".Length; ++sio)\n");
				sb.Append(s_indent + "{\n");
				sb.Append(s_indent2 + "for(int sim = 0; sim < " + sourcergn + "[sio]." + cards + "; ++sim)\n");
				sb.Append(s_indent2 + "{\n");

				sb.Append(s_indent3 + "int pc = PexChoose.ValueFromRange(\"trgtrgn\", 0, " + (trgtrgns.Count + (maynull ? 1 : 0)) + ");\n");
				sb.Append(s_indent3 + "int pcv;\n");

				sb.Append(s_indent3 + "switch(pc)\n");
				sb.Append(s_indent3 + "{\n");
				for(int j = 0; j < trgtrgns.Count; ++j)
				{
					sb.Append(s_indent3 + "case " + j + ":\n");
					sb.Append(s_indent4 + s_assume + "(" + trgtrgns[j] + ".Length != 0);\n");
					sb.Append(s_indent4 + "pcv = PexChoose.ValueFromRange(\"trgtobj\", 0, " + trgtrgns[j] + ".Length);\n");

					if(!mayintf)
						sb.Append(s_indent4 + s_assume + "(" + trgtrgns[j] + "[pcv] != null);\n");

					sb.Append(s_indent4 + sourcergn + "[sio][sim]" + " = " + trgtrgns[j] + "[pcv];\n");

					if(!mayintf)
						sb.Append(s_indent4 + trgtrgns[j] + "[pcv] = null;\n");

					sb.Append(s_indent4 + "break;\n");
				}

				if(maynull)
				{
					sb.Append(s_indent3 + "default:\n");
					sb.Append(s_indent4 + sourcergn + "[sio][sim]" + " = null;\n");
					sb.Append(s_indent4 + "break;\n");
				}

				sb.Append(s_indent3 + "}\n");
				sb.Append(s_indent2 + "}\n");
				sb.Append(s_indent + "}\n");
			}
			else
			{
				string mfname = mfid.FieldName;
				IRType ftype = (mfid.FieldType is IRPtrType ? ((IRPtrType)mfid.FieldType).PtrTargetType : mfid.FieldType);

				sb.Append(s_indent + "for(int k = 0; k < " + sourcergn + ".Length; ++k)\n");
				sb.Append(s_indent + "{\n");

				if(trgtrgns.Count == 1 && (!maynull))
				{
					sb.Append(s_indent2 + "int pcv = PexChoose.ValueFromRange(\"trgtobj\", 0, " + trgtrgns[0] + ".Length);\n");

					if(!mayintf)
						sb.Append(s_indent2 + s_assume + "(" + trgtrgns[0] + "[pcv] != null);\n");

					sb.Append(s_indent2 + "PexInvariant.SetField<" + ftype.ToString() + ">(");
					sb.Append(sourcergn + "[k]," + " \"" + mfname + "\", " + trgtrgns[0] + "[pcv]);\n");

					if(!mayintf)
						sb.Append(s_indent2 + trgtrgns[0] + "[pcv] = null;\n");
				}
				else
				{
					sb.Append(s_indent2 + s_assume + "(" + trgtrgns[0] + ".Length != 0);\n");
					sb.Append(s_indent2 + "int pc = PexChoose.ValueFromRange(\"trgtrgn\", 0, " + (trgtrgns.Count + (maynull ? 1 : 0)) + ");\n");
					sb.Append(s_indent2 + "int pcv;\n");

					sb.Append(s_indent2 + "switch(pc)\n");
					sb.Append(s_indent2 + "{\n");
					for(int j = 0; j < trgtrgns.Count; ++j)
					{
						sb.Append(s_indent2 + "case " + j + ":\n");
						sb.Append(s_indent3 + s_assume + "(" + trgtrgns[j] + ".Length != 0);\n");
						sb.Append(s_indent3 + "pcv = PexChoose.ValueFromRange(\"trgtobj\", 0, " + trgtrgns[j] + ".Length);\n");

						if(!mayintf)
							sb.Append(s_indent3 + s_assume + "(" + trgtrgns[j] + "[pcv] != null);\n");

						sb.Append(s_indent3 + "PexInvariant.SetField<" + ftype.ToString() + ">(");
						sb.Append(sourcergn + "[k]," + " \"" + mfname + "\", " + trgtrgns[0] + "[pcv]);\n");

						if(!mayintf)
							sb.Append(s_indent3 + trgtrgns[j] + "[pcv] = null;\n");

						sb.Append(s_indent3 + "break;\n");
					}

					if(maynull)
					{
						sb.Append(s_indent2 + "default:\n");
						sb.Append(s_indent3 + "PexInvariant.SetField<" + ftype.ToString() + ">(");
						sb.Append(sourcergn + "[k]," + "\"" + mfname + "\"," + "null);\n");
						sb.Append(s_indent3 + "break;\n");
					}

					sb.Append(s_indent2 + "}\n");
				}

				sb.Append(s_indent + "}\n");
			}

			return sb.ToString();
		}

		/// <summary>
		/// Generate a string which goes over the scalar fields in the objects and asks pex for a value
		/// </summary>
		private string GenInitScalars()
		{
			StringBuilder csb = new StringBuilder();

			foreach(PIMNode hr in this.m_pim.Nodes)
			{
				if(hr.HomeLocFor != null)
					continue;

				string rgntypect = NameGen.GenerateNameForRegionCount(hr);
				string rgntypename = NameGen.GenerateNameForRegion(hr);

				StringBuilder tcsb = new StringBuilder();

				if(hr.LocationType is IRContainerType)
				{
					if(!(((IRContainerType)hr.LocationType).ContainerContentsType is IRPtrType))
					{
						IRType cctype = ((IRContainerType)hr.LocationType).ContainerContentsType;

						tcsb.Append(s_indent + "for(int iilv = 0; iilv < " + rgntypect + "; ++iilv)\n");
						tcsb.Append(s_indent + "{\n");

						tcsb.Append(s_indent2 + "for(int ii = 0; ii < " + rgntypename + "[iilv]." + (((IRContainerType)hr.LocationType).IsArrayContainer ? "Length" : "Count") + "; ++ii)\n");
						if(cctype.TypeName.Equals("System.Bool"))
							tcsb.Append(s_indent3 + rgntypename + "[iilv][ii] = PexChoose.Value<bool>(\"bv\");\n");
						else if(cctype.TypeName.Equals("System.Int32"))
							tcsb.Append(s_indent3 + rgntypename + "[iilv][ii] = PexChoose.Value<int>(\"iv\");\n");
						else if(cctype.TypeName.Equals("System.Double"))
							tcsb.Append(s_indent3 + rgntypename + "[iilv][ii] = PexChoose.Value<double>(\"dv\");\n");
						else
						{ ;}

						tcsb.Append(s_indent + "}\n");
					}
				}
				else
				{
					List<IRBasicAccessPath> mfid = hr.LocationType.TypeDescription.AllTerminalAccessPaths;
					foreach(IRBasicAccessPath bap in mfid)
					{
						string mfname = bap.PathLast.FieldName;
						IRType ftype = bap.PathLast.FieldType;
						if(ftype is IRPtrType)
							continue;

						if(bap.PathLast.FieldName.Equals("rtti"))
							continue;

						tcsb.Append(s_indent + "for(int ii = 0; ii < " + rgntypect + "; ++ii)\n");

						if(ftype.TypeName.Equals("System.Bool"))
							tcsb.Append(s_indent2 + "PexInvariant.SetField<bool>(" + rgntypename + "[ii]," + " \"" + mfname + "\"," + " PexChoose.Value<bool>(\"bv\"));\n");
						else if(ftype.TypeName.Equals("System.Int32"))
							tcsb.Append(s_indent2 + "PexInvariant.SetField<int>(" + rgntypename + "[ii]," + " \"" + mfname + "\"," + " PexChoose.Value<int>(\"iv\"));\n");
						else if(ftype.TypeName.Equals("System.Double"))
							tcsb.Append(s_indent2 + "PexInvariant.SetField<double>(" + rgntypename + "[ii]," + " \"" + mfname + "\"," + " PexChoose.Value<double>(\"dv\"));\n");
						else
						{ ;}
					}
				}

				if(tcsb.Length != 0)
					csb.Append(tcsb + "\n");
			}

			return csb.ToString();
		}

		private String BuildConstructorArgList(string oparams, List<string> paramargs)
		{
			StringBuilder alsb = new StringBuilder(oparams);

			if(paramargs.Count != 0 && (!String.IsNullOrEmpty(oparams)))
				alsb.Append(", ");

			for(int i = 0; i < paramargs.Count; ++i)
			{
				if(i != 0) alsb.Append(", ");
				alsb.Append("int " + paramargs[i]);
			}

			return "(" + alsb.ToString() + ")\n";
		}
	}
}
